import pandas as pd
from sklearn.preprocessing import OneHotEncoder
from sklearn.feature_selection import chi2
from pandas.core.frame import DataFrame
import numpy as np


def get_data(path, nans=True):
    data_source = pd.read_csv(path)
    if nans == False:
        return data_source.dropna(axis='index')
    return data_source


def data_pretreat(path):  # 定义数值替换文本的函数
    data = get_data(path, nans=False)
    data = data.drop(columns='number')
    # 将源文件中3种协议类型转换成数字标识
    data['protocol_type'] = data['protocol_type'].map({'tcp': 0, 'udp': 1, 'icmp': 2})
    # 将源文件中70种网络服务类型转换成数字标识
    data['service'] = data['service'].map(
        {'aol': 0, 'auth': 1, 'bgp': 2, 'courier': 3, 'csnet_ns': 4, 'ctf': 5, 'daytime': 6, 'discard': 7, 'domain': 8,
         'domain_u': 9, 'echo': 10, 'eco_i': 11, 'ecr_i': 12, 'efs': 13, 'exec': 14, 'finger': 15, 'ftp': 16,
         'ftp_data': 17, 'gopher': 18, 'harvest': 19, 'hostnames': 20, 'http': 21, 'http_2784': 22, 'http_443': 23,
         'http_8001': 24, 'imap4': 25, 'IRC': 26, 'iso_tsap': 27, 'klogin': 28, 'kshell': 29, 'ldap': 30, 'link': 31,
         'login': 32, 'mtp': 33, 'name': 34, 'netbios_dgm': 35, 'netbios_ns': 36, 'netbios_ssn': 37, 'netstat': 38,
         'nnsp': 39, 'nntp': 40, 'ntp_u': 41, 'other': 42, 'pm_dump': 43, 'pop_2': 44, 'pop_3': 45, 'printer': 46,
         'private': 47, 'red_i': 48, 'remote_job': 49, 'rje': 50, 'shell': 51, 'smtp': 52, 'sql_net': 53, 'ssh': 54,
         'sunrpc': 55, 'supdup': 56, 'systat': 57, 'telnet': 58, 'tftp_u': 59, 'tim_i': 60, 'time': 61, 'urh_i': 62,
         'urp_i': 63, 'uucp': 64, 'uucp_path': 65, 'vmnet': 66, 'whois': 67, 'X11': 68, 'Z39_50': 69})
    # 将源文件中11种网络连接状态转换成数字标识
    data['flag'] = data['flag'].map(
        {'OTH': 0, 'REJ': 0, 'RSTO': 0, 'RSTOS0': 0, 'RSTR': 0, 'S0': 0, 'S1': 0, 'S2': 0, 'S3': 0, 'SF': 1, 'SH': 0})
    # 将源文件中攻击类型转换成数字标识(训练集中共出现了22个攻击类型，而剩下的17种只在测试集中出现)
    data['label'] = data['label'].map(
        {'normal': 0, 'ipsweep': 1, 'mscan': 2, 'nmap': 3, 'portsweep': 4, 'saint': 5, 'satan': 6, 'apache2': 7,
         'back': 8, 'land': 9, 'mailbomb': 10, 'neptune': 11, 'pod': 12, 'processtable': 13, 'smurf': 14,
         'teardrop': 15, 'udpstorm': 16, 'buffer_overflow': 17, 'httptunnel': 18, 'loadmodule': 19, 'perl': 20,
         'ps': 21, 'rootkit': 22, 'sqlattack': 23, 'xterm': 24, 'ftp_write': 25, 'guess_passwd': 26, 'imap': 27,
         'multihop': 28, 'named': 29, 'phf': 30, 'sendmail': 31, 'snmpgetattack': 32, 'snmpguess': 33, 'spy': 34,
         'warezclient': 35, 'warezmaster': 36, 'worm': 37, 'xlock': 38, 'xsnoop': 39})
    # 数值归一化:最值归一化
    '''data['service'] = (data['service'] - data['service'].min()) / (data['service'].max() - data['service'].min())
    data['src_bytes'] = (data['src_bytes'] - data['src_bytes'].min()) / (
            data['src_bytes'].max() - data['src_bytes'].min())
    data['dst_bytes'] = (data['dst_bytes'] - data['dst_bytes'].min()) / (
            data['dst_bytes'].max() - data['dst_bytes'].min())
    data['count'] = (data['count'] - data['count'].min()) / (data['count'].max() - data['count'].min())
    data['srv_count'] = (data['srv_count'] - data['srv_count'].min()) / (
            data['srv_count'].max() - data['srv_count'].min())
    data['dst_host_count'] = (data['dst_host_count'] - data['dst_host_count'].min()) / (
            data['dst_host_count'].max() - data['dst_host_count'].min())
    data['dst_host_srv_count'] = (data['dst_host_srv_count'] - data['dst_host_srv_count'].min()) / (
            data['dst_host_srv_count'].max() - data['dst_host_srv_count'].min())'''
    np.isnan(data).any()  # data里是否存在nan
    data.dropna(inplace=True)  # 删除有缺失值的行
    return data


if __name__ == '__main__':
    data = data_pretreat("../resource/KDDTrain+.csv")
    print(data.head())
    X = data.loc[:, data.columns != 'label']
    Y = data['label']
    print(Y.value_counts())
    # 卡方检验 https://www.jianshu.com/p/64974c4de9d4
    chi_scores = chi2(X, Y)
    print(chi_scores)
    p_values = pd.Series(data=chi_scores[0] + chi_scores[1], index=X.columns)
    p_values.sort_values(ascending=True, inplace=True)
    print('Values in order of ascending p-values (lower=more significant)')
    print(p_values)
